// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

namespace Microsoft.Quantum.QsCompiler.Experimental

open System
open System.Collections.Immutable
open System.Text.RegularExpressions
open Microsoft.Quantum.QsCompiler.DataTypes
open Microsoft.Quantum.QsCompiler.Experimental.Utils
open Microsoft.Quantum.QsCompiler.SyntaxTokens
open Microsoft.Quantum.QsCompiler.SyntaxTree
open Microsoft.Quantum.QsCompiler.Transformations.Core


/// The ScopeTransformation used to ensure unique variable names.
/// When called on a function body, will transform it such that all local variables defined
/// in the function body have unique names, generating new variable names if needed.
/// Autogenerated variable names have the form __qsVar[X]__[name]__.
type VariableRenaming() =
    inherit SyntaxTreeTransformation()

    /// Returns a copy of the given variable stack inside of a new scope
    let enterScope map = Map.empty :: map

    /// Returns a copy of the given variable stack outside of the current scope.
    /// Throws an ArgumentException if the given variable stack is empty.
    let exitScope = List.tail

    /// Returns the value associated to the given key in the given variable stack.
    /// If the key is associated with multiple values, returns the one highest on the stack.
    /// Returns None if the key isn't associated with any values.
    let tryGet key = List.tryPick (Map.tryFind key)

    /// Returns a copy of the given variable stack with the given key set to the given value.
    /// Throws an ArgumentException if the given variable stack is empty.
    let set (key, value) = function
        | [] -> ArgumentException "No scope to define variables in" |> raise
        | head :: tail -> Map.add key value head :: tail

    /// A regex that matches the original name of a mangled variable name
    let varNameRegex = Regex("^__qsVar\d+__(.+)__$")


    /// The number of times a variable is referenced
    let mutable newNamesSet = Set.empty
    /// The current dictionary of new names to substitute for variables
    let mutable renamingStack = [Map.empty]
    /// Whether we should skip entering the next scope we encounter
    let mutable skipScope = false

    /// Given a possibly-mangled variable name, returns the original variable name
    let demangle varName =
        let m = varNameRegex.Match varName
        if m.Success then m.Groups.[1].Value else varName

    /// Given a new variable name, generates a new unique name and updates the state accordingly
    let generateUniqueName varName =
        let baseName = demangle varName
        let mutable num, newName = 0, baseName
        while newNamesSet.Contains newName do
            num <- num + 1
            newName <- sprintf "__qsVar%d__%s__" num baseName
        newNamesSet <- newNamesSet.Add newName
        renamingStack <- set (varName, newName) renamingStack
        newName

    /// Processes the initial argument tuple from the function declaration
    let rec processArgTuple = function
        | QsTupleItem {VariableName = ValidName name} -> generateUniqueName name.Value |> ignore
        | QsTupleItem {VariableName = InvalidName} -> ()
        | QsTuple items -> Seq.iter processArgTuple items

    member __.clearStack() =
        renamingStack <- [Map.empty]


    override __.onProvidedImplementation (argTuple, body) =
        newNamesSet <- Set.empty
        renamingStack <- [Map.empty]
        do processArgTuple argTuple
        base.onProvidedImplementation (argTuple, body)

    override __.Scope = { new ScopeTransformation() with

        override __.Transform x =
            if skipScope then
                skipScope <- false
                base.Transform x
            else
                renamingStack <- enterScope renamingStack
                let result = base.Transform x
                renamingStack <- exitScope renamingStack
                result

        override __.Expression = { new ExpressionTransformation() with
            override expr.Kind = { new ExpressionKindTransformation() with
                override __.ExpressionTransformation ex = expr.Transform ex
                override __.TypeTransformation t = t

                override __.onIdentifier (sym, tArgs) =
                    maybe {
                        let! name =
                            match sym with
                            | LocalVariable name -> Some name.Value
                            | _ -> None
                        let! newName = tryGet name renamingStack
                        return Identifier (LocalVariable (NonNullable<_>.New newName), tArgs)
                    } |? Identifier (sym, tArgs)
            }
        }

        override this.StatementKind = { new StatementKindTransformation() with
            override __.ExpressionTransformation x = this.Expression.Transform x
            override __.LocationTransformation x = x
            override __.ScopeTransformation x = this.Transform x
            override __.TypeTransformation x = x

            override this.onSymbolTuple syms =
                match syms with
                | VariableName item -> VariableName (NonNullable<_>.New (generateUniqueName item.Value))
                | VariableNameTuple items -> Seq.map this.onSymbolTuple items |> ImmutableArray.CreateRange |> VariableNameTuple
                | InvalidItem | DiscardedItem -> syms

            override __.onRepeatStatement stm =
                renamingStack <- enterScope renamingStack
                skipScope <- true
                let result = base.onRepeatStatement stm
                renamingStack <- exitScope renamingStack
                result
        }
    }
